#!usr/bin/env/ python
# -*- coding:utf-8 -*-
"""
Author:Xiaoxu Zhang
Date:2024-05-21
"""
from kan import KAN
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons
import torch

dataset = {}
train_input, train_label = make_moons(n_samples=1000, shuffle=True, noise=0.1, random_state=None)
test_input, test_label = make_moons(n_samples=1000, shuffle=True, noise=0.1, random_state=None)

dataset['train_input'] = torch.from_numpy(train_input)
dataset['test_input'] = torch.from_numpy(test_input)
dataset['train_label'] = torch.from_numpy(train_label[:,None])
dataset['test_label'] = torch.from_numpy(test_label[:,None])

X = dataset['train_input']
y = dataset['train_label']
plt.scatter(X[:, 0], X[:, 1], c=y[:, 0])
plt.savefig('../out/moon-like.png')


model = KAN(width=[2, 1], grid=3, k=3)


def train_acc():
    return torch.mean((torch.round(model(dataset['train_input'])[:, 0]) == dataset['train_label'][:, 0]).float())


def test_acc():
    return torch.mean((torch.round(model(dataset['test_input'])[:,0]) == dataset['test_label'][:,0]).float())


results = model.train(dataset, opt="LBFGS", steps=20, metrics=(train_acc, test_acc))

# plot loss data
plt.cla()
plt.plot(results['train_loss'], label='train loss')
plt.plot(results['test_loss'], label='test loss')
plt.legend()
plt.savefig('../out/kan-loss.png')

# plot acc data
plt.cla()
plt.plot(results['train_acc'], label='train_acc')
plt.plot(results['test_acc'], label='test_acc')
plt.legend()
plt.title("train by kan")
plt.savefig('../out/kan-acc.png')

# plot reg data
plt.plot(results['reg'], label='reg')
plt.savefig('../out/kan-reg.png')
